# -*- coding: utf-8 -*-
import os
import requests
import shutil
from shutil import copyfile
import time
from Crypto.Cipher import AES    #注：python3 安装 Crypto 是 pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple pycryptodome

# 创建目录
def mkDir(path):
    """
    创建目录
    :param path: 目录
    :return:
    """
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \\ / 符号
    path = path.rstrip("\\")
    path = path.rstrip("/")
    # 判断路径是否存在
    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        os.makedirs(path)
    return True

# 创建一个文件并写入内容
def mkFile(path,text):
    """
    创建一个文件并写入内容
    :param path: E:/file/a.txt
    :param text: 我是一个文件的内容
    :return: 成功状态
    """
    f = open(path, "w")
    f.write(text)
    f.close()
    return True

# 删除某个目录下的所有文件夹和文件
def killdirData(path):
    """
    删除某个目录下的所有文件夹和文件
    :param path: E:mp4/demo
    """
    ls = os.listdir(path)
    for i in ls:
        c_path = os.path.join(path, i)
        if os.path.isdir(c_path):
            shutil.rmtree(c_path)
        else:
            os.remove(c_path)

# 删除指定目录下的某个文件
def removeFile(path):
    """
    删除指定目录下的某个文件
    :param path: E:mp4/a.txt
    :return:
    """
    if os.path.exists(path):
        os.remove(path)

# 复制某个文件到另一个目录
def copyFile(path1,path2):
    """
    复制某个文件到另一个目录
    :param path1: E:mp4/a.txt
    :param path2: D:mp3/jay/a.txt
    :return: 成功状态
    """
    copyfile(path1, path2)

# 修改文件名
def exitFileName(path, name):
    """
    修改文件名
    :param path: E:mp4/a.txt
    :param name: b.txt
    :return:
    """
    path = path.replace("\\", "/")
    path2 = path.split(path.split("/")[-1])[0]+"/"
    os.rename(path, path2 + name)

# 检测m3u8文件是否需要解密
def fileIfDecrypt(path):
    """
    检测m3u8文件是否需要解密
    :param path: E:mp4/a.m3u8
    :return: 需要解密返回key路径，不需要解密就返回False
    """
    with open(path) as f:
        for line in f:
            if "EXT-X-KEY" in line:
                key_path = line[line.find("URI"):line.rfind('"')].split('"')[1]
                f.close()
                if key_path[0] is "/":
                    return key_path
                else:
                    return "/" + key_path
    f.close()
    return False

# 下载任何格式的网络文件
def downloadFile(url,filepath):
    """
    下载任何格式的网络文件
    :param url: 文件的链接地址   http://cn5.download05.com/hls/name.txt
    :param filepath: 存放的本地路径包含文件名 E:mp4/a.txt
    :return: 返回成功状态和耗时(秒)
    """
    if os.path.exists(filepath):
        return "True,0.300"
    else:
        try:
            header = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
                                    " (KHTML, like Gecko) Chrome/74.0.3729.108 Safari/537.36"}

            st = int(round(time.time() * 1000))
            r = requests.get(url, headers=header, timeout=15)
            r.raise_for_status()
            with open(filepath, "wb") as f:  # wb：以二进制方式写入文件，如果文件存在会覆盖
                f.write(r.content)  # r.content：以二进制方式读取文件,然后写入
            f.close()
            et = int(round(time.time() * 1000))
            ht = str((et-st)/1000)  # 耗时
            return "True," + ht
        except:
            return "False,0.300"

# 获取秘钥
def getKey(m3u8url,keyurl,keypath):
    """
    获取秘钥
    :param m3u8url: m3u8有效连接地址 http://cn5.download05.com/hls/20190525/index.m3u8
    :param keyurl:  /20190525/GPlWuHfM/1561kb/hls/key.key
    :param keypath: E:mp4/a
    :return: 秘钥字符串
    """
    zt = False
    print("正在获取秘钥...")
    list = []
    list.append(m3u8url.split("//")[0]+"//"+m3u8url.split("//")[1].split("/")[0])
    list.append(m3u8url[0:m3u8url.rfind("/")])

    keypath = keypath+"/"+keyurl[keyurl.rfind("/"):]

    for u1 in list:
        ulist = u1.split("/")
        klist = keyurl.split("/")
        for k1 in klist:
            if (k1 in ulist) and k1 != '':
                u1 = u1.split(k1)[0]
        url = u1 + keyurl

        removeFile(keypath)  # 下载前先删除，确保此文件不存在，防止文件失效而一直下载不到
        t1 = downloadFile(url, keypath)
        if "True" in t1:
            f = open(keypath, "r")
            text = f.read().strip("\n").strip()
            f.close()
            zt = True
            print("秘钥获取成功: " + text)
            return text
    return zt

# 下载ts文件并解密
def downloadTsFileAndDecrypt(url,filepath,key):
    """
    下载ts文件并解密
    :param url:  https://3atvplay.com/20190801/HRraP1tT/1259kb/hls/asd123.ts
    :param path: E:/mp4/a.ts
    :param key:  508630799be9fc8a
    :return:     成功状态
    """
    if os.path.exists(filepath):
        return "True,0.300"
    else:
        try:
            header = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
                                    " (KHTML, like Gecko) Chrome/74.0.3729.108 Safari/537.36"}

            key = bytes(key, encoding="utf8")
            cryptor = AES.new(key, AES.MODE_CBC, key)

            st = int(round(time.time() * 1000))
            r = requests.get(url, headers=header, timeout=15)
            r.raise_for_status()
            with open(filepath, "wb") as f:  # wb：以二进制方式写入文件，如果文件存在会覆盖
                f.write(cryptor.decrypt(r.content))
            f.close()
            et = int(round(time.time() * 1000))
            ht = str((et-st)/1000)  # 耗时
            return "True," + ht
        except:
            return "False,0.300"

# 解密本地ts文件
def decryptLocalTsFile(path,key):
    """
    解密本地ts文件
    :param path: E:\\mp4\\a.ts
    :param key:  508630799be9fc8a
    :return:     状态
    """
    try:
        # 在当前目录下创建一个新文件，存放解析后的新文件
        path = path.replace("\\","/")
        path2 = path.split(path.split("/")[-1])[0] + "jiexi"
        newfilename = path2 + "/" + path.split("/")[-1]
        mkDir(path2)

        key = bytes(key, encoding="utf8")
        cryptor = AES.new(key, AES.MODE_CBC, key)
        a = open(path, "rb")
        with open(newfilename, "wb") as f:  # wb：以二进制方式写入文件，如果文件存在会覆盖
            f.write(cryptor.decrypt(a.read()))
        a.close()
        f.close()
        print("解密成功 "+path)
        return True
    except:
        print("解密文件失败 " +path)
        return False

# 下载m3u8文件到指定目录，返回可直接拼接ts文件名的有效链接，m3u8存放路径
def downM3u8(url,path):
    """
    根据链接地址下载m3u8文件到指定目录
    :param url: http://cn5.download05.com/hls/20190812/182d4d695b9050c77f140bc0b21fbe9e/1565606686/index.m3u8
    :param path: E:mp4/hongfanqu/index.m3u8  （如果目录不存在不会新建目录）
    :return: 可直接拼接ts文件名的有效链接，m3u8存放路径
    """
    print("正在下载m3u8文件，请稍后...")
    t = downloadFile(url, path)
    if "True" in t:
        # ts文件链接地址拼接方案
        list = []
        list.append(url.split("/")[0] + "//" + url.split("/")[1] + url.split("/")[2])
        list.append(url.split("/index.m3u8")[0])

        iscg = True
        for url2 in list:
            with open(path) as f:
                for line in f:
                    if "index.m3u8" in line:
                        if line[0] is "/":
                            url3 = url2 + (line.split('\n')[0])
                        else:
                            url3 = url2 + "/" + (line.split('\n')[0])
                        path2 = path.split("index.m3u8")[0] + "index"
                        mkDir(path2)
                        t2 = downloadFile(url3, path2 + "/index.m3u8")
                        if "True" in t2:
                            print(path2 + "/index.m3u8 文件下载成功")
                            f.close()
                            return url3.split("/index.m3u8")[0] + "," + path2
                        else:
                            iscg = False
                            print("m3u8文件下载失败，正在重新组装链接...")
                if iscg is True:
                    break
        print(path + " 文件下载成功")
        f.close()
        return url.split("index.m3u8")[0] + "," + path
    else:
        print(path + " 文件下载失败")
        return False

# 解析m3u8文件返回ts链接地址 （列表）
def analysisM3u8File(path,url):
    """
    解析m3u8文件返回ts链接地址 （列表）
    :param path: m3u8文件地址 E:mp4/hongfanqu/index.m3u8
    :param url:  有效的头链接  http://cn5.download05.com/100kb/hls
    :return: ts链接列表
    """
    print("正在解析m3u8文件： " + path)
    urls = []
    with open(path, "r") as file:
        for line in file:
            if line.endswith(".ts\n"):
                if line[:1] is "/":
                    urls.append(url + line.strip("\n"))
                else:
                    urls.append(url + "/" + line.strip("\n"))
    file.close()
    print("解析完毕，ts链接地址已生成")
    return urls

# 根据ts链接地址，下载部分ts文件
def downSomeTs(url,st,en,path):
    """
    根据ts链接地址，下载部分ts文件
    :param url: ts文件链接地址
    :param st: ts文件开始编号
    :param en: ts文件结束编号
    :param path: 本地存放路径
    :return: 成功状态
    """

    print("正在下载ts文件，请稍后...")

    mkDir(path)
    st = int(st)
    en = int(en)
    url = url.strip()

    # urlf = open(path + "/url.txt", "w")
    # urlf.write(url)
    # urlf.close()

    u2 = url.split(".ts")[0].split("/")[-1]
    url1 = url.split(u2)[0]

    # https://dadi-yun.com/20190615/9821_64e278a8/800k/hls/9yEN6OZ8959998.ts

    length = en-st+1  # 文件个数
    ishb = True
    num = 0
    for i in range(length):
        s1 = int(u2[-6:-4] + (6-len(u2[-6:-4]))*"0")
        num = num+1
        path2 = u2[:-6] + str(s1 + st) + ".ts"
        http = url1+path2
        t = downloadFile(http, path+"/"+path2)
        if "True" in t:
            sy = length - num  # 剩余多少个
            htime = round((float(t.split(",")[1]) * sy) / 60, 2)
            print(path + path2 + " 下载成功, 总共"+str(length)+"个，剩余"+str(sy)+"个, 预计需要 "+str(htime)+" 分钟")
        else:
            ishb = False
            print(path + path2 + " 下载失败")
        st = st + 1

    if ishb is True:
        return hbFile(path)
    else:
        print("ts文件残缺，无法开始合并")
        return False

# 合并文件
def hbFile(localpath):
    """
    合并目录下所有的ts文件
    :param localpath:
    :return: new.ts
    """
    print(localpath+" 正在合并文件")
    try:
        newfile = open(localpath + "/" + str(int(time.time()))+".ts", mode="ab")
        a = os.listdir(localpath)
        for i in a:
            if "ts" in i:
                f = open(localpath + "/" + i, mode="rb")
                newfile.write(f.read())
                f.close()
        newfile.close()
        print(localpath + " 目录文件合并成功")
        return True
    except:
        print(localpath+" 合并失败")
        return False


def run(url, localpath):
    """
    根据m3u8链接下载ts文件
    :param url: m3u8文件链接地址
    :param localpath: 存在的路径 （目录）
    :return: 成功状态
    """
    url = url.strip()
    now_time = str(int(time.time()))
    localpath += "/" + now_time  # 设定子文件夹
    mkDir(localpath)  # 创建子文件夹

    t = downM3u8(url, localpath + "/index.m3u8")
    if t is not False:
        localpath = t.split(",")[1]

        urlf = open(localpath + "/url.txt", "w")
        urlf.write(url)
        urlf.close()


        print("m3u8文件下载成功，正在解析")
        url2 = t.split(",")[0]
        list = []
        with open(localpath + "/index.m3u8", "r") as fr:
            for line in fr:  # 一行一行读取
                if ".ts" in line:
                    if line[:1] is "/":
                        ts_urlpath = url2 + (line.split('\n')[0])
                    else:
                        ts_urlpath = url2.split("index")[0] + "/" + (line.split('\n')[0])
                    list.append(ts_urlpath)
        print("m3u8文件解析成功，开始下载ts文件")
        num = 0
        length = len(list)
        for i in list:
            tsfilename = localpath + "/" + i.split("/")[-1]
            t2 = downloadFile(i, tsfilename)
            num = num + 1
            if "True" in t2:
                sy = length - num  # 剩余多少个
                htime = round(float(t2.split(",")[1])*sy/60, 2)
                print(tsfilename + "下载成功，总共" + str(length) + "个，剩余" + str(sy) + "个, 预计需要 "+str(htime)+" 分钟")
            else:
                print(tsfilename + "下载失败")
        print("ts文件全部下载完成，开始合并ts文件")
        return True
    else:
        print("m3u8文件下载失败")
        return False

def run2(url,path):
    """
    根据本地m3u8文件下载ts文件
    :param url: m3u8链接地址
    :param path: ts文件存放目录
    :return:
    """
    url = url.strip()
    if "/index.m3u8" in url:
        url = url.split("/index.m3u8")[0]

    list = []
    with open(path + "/index.m3u8", "r") as fr:
        for line in fr:  # 一行一行读取
            if ".ts" in line:
                if line[:1] is "/":
                    ts_urlpath = url + (line.split('\n')[0])
                else:
                    ts_urlpath = url + "/" + (line.split('\n')[0])
                list.append(ts_urlpath)
    print("m3u8文件解析成功，开始下载ts文件")
    num = 0
    length = len(list)
    for i in list:
        tsfilename = path + "/" + i.split("/")[-1]
        t2 = downloadFile(i, tsfilename)
        num = num + 1
        if "True" in t2:
            sy = length - num  # 剩余多少个
            htime = round(float(t2.split(",")[1]) * sy / 60, 2)
            print(tsfilename + "下载成功，总共" + str(length) + "个，剩余" + str(length - num) + "个,预计需要"+str(htime)+"分钟")
        else:
            print(i + " 下载失败")
    print("ts文件全部下载完成，开始合并ts文件")
    return True



# getKey("https://3atvplay.com/20190525/GPlWuHfM/1561kb/hls/index.m3u8","/20190525/GPlWuHfM/1561kb/hls/key.key","E:mp4/1")
# downloadFile("https://3atvplay.com/20190801/HRraP1tT/1259kb/hls/GPddVl7842004.ts","E:/mp4/a.ts")
# d3 = downM3u8("https://3atvplay.com/20190525/GPlWuHfM/index.m3u8","E:/mp4/1/index.m3u8")
# print(d3)
# downSomeTs("https://dadi-yun.com/20190614/9802_73605560/800k/hls/9yEN6OZ8959998.ts.ts","2345","2400","E:/mp4/1")
# downloadTsFileAndDecrypt("https://3atvplay.com/20190801/HRraP1tT/1259kb/hls/GPddVl7842004.ts",
#                          "E:/mp4/a.ts","508630799be9fc8a")
# decryptLocalTsFile("E:\\mp4\\a.ts","508630799be9fc8a")

print("******************欢迎使用小斌ts文件下载软件*******************")
while True:
    print("**--------------软件说明---------------**")
    print("1. 文件默认存放目录为 E:/mp4")
    print("   每下载一个m3u8文件都会在此目录下单独创建子文件")

    print("**------------------------------------**")
    print("请选择ts文件获取方式")
    print("     0 网络链接下载m3u8文件（一般用于创建一个新任务）")
    print("     1 指定本地m3u8文件路径（一般用于继续上次未完成的任务）")
    print("     2 直接下载ts文件并合并（一般用于直接下载视频中的某一段）")
    # print("         提示：ts文件名后6位必须是数字")
    print("     3 合并ts文件")
    print("     4 退出")

    fs = input(">>请输入指令号: ")
    if fs is "0":
        url = input("请输入m3u8地址: ")
        localpath = "E:/mp4"
        t1 = run(url, localpath)  # 下载ts文件
        if t1 is True:
            hb = hbFile(localpath)  # 合并ts文件
            if hb is True:
                print("任务完成...")
    elif fs is "1":
        url = input("请输入m3u8地址: ")
        u8path = input("请指定m3u8文件所在目录: ")
        t2 = run2(url, u8path)
        if t2 is True:
            hb = hbFile(u8path)  # 合并ts文件
            if hb is True:
                print("任务完成...")
    elif fs is "2":
        url = input("请输入ts文件链接地址: ")
        st = input("请输入ts文件开始号(后4位): ")
        ed = input("请输入ts文件结束号(后4位): ")

        if int(ed) > 9999 or int(st) > 9999:
            print("暂时不支持文件号大于9999的文件,请使用指令0完成任务")
            continue
        elif int(ed) < int(st):
            print("结束号不能小于开始号,请使用指令0完成任务")
            continue

        path = input("请输入存放路径: ")
        while True:
            t3 = downSomeTs(url, st, ed, path)
            if t3 is True:
                print("任务完成...")
                break
            else:
                va = input("是否重新开始？（y/n）")
                if va is "n":
                    break
    elif fs is "3":
        path = input("请输入需要合并的路径: ")
        hb = hbFile(path)
        if hb is True:
            print("任务完成...")
    else:
        break



